#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
数据库工具模块

提供数据库连接池、上下文管理器和辅助函数
"""

import sqlite3
import threading
import queue
import time
from contextlib import contextmanager
from typing import Dict, Any, List, Optional, Tuple
from flask import current_app


class SQLiteConnectionPool:
    """SQLite连接池，使用单例模式确保全局唯一"""
    
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls, *args, **kwargs):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(SQLiteConnectionPool, cls).__new__(cls)
                cls._instance._initialized = False
            return cls._instance
    
    def __init__(self, db_path: str = None, max_size: int = 10, timeout: int = 30):
        """初始化连接池
        
        Args:
            db_path: 数据库文件路径
            max_size: 连接池最大连接数
            timeout: 获取连接超时时间(秒)
        """
        # 避免重复初始化
        if self._initialized:
            return
            
        self.db_path = db_path
        self.max_size = max_size
        self.timeout = timeout
        self._pool = queue.Queue(maxsize=max_size)
        self._size = 0
        self._lock = threading.RLock()
        self._initialized = True
        
        # 预创建连接
        self._create_initial_connections()
    
    def _create_initial_connections(self, initial_size: int = 5):
        """预创建初始连接
        
        Args:
            initial_size: 初始连接数量
        """
        for _ in range(min(initial_size, self.max_size)):
            self._add_connection()
    
    def _add_connection(self):
        """创建并添加一个新连接到池中"""
        if self._size >= self.max_size:
            return False
            
        with self._lock:
            if self._size >= self.max_size:
                return False
                
            conn = sqlite3.connect(
                self.db_path,
                detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES,
                check_same_thread=False  # 允许在不同线程中使用
            )
            # 设置行工厂为字典，方便处理结果
            conn.row_factory = sqlite3.Row
            
            # 启用外键支持
            conn.execute("PRAGMA foreign_keys = ON")
            
            # 添加到池中
            self._pool.put(conn)
            self._size += 1
            return True
    
    def get_connection(self):
        """从连接池获取一个连接
        
        Returns:
            sqlite3.Connection: 数据库连接
        
        Raises:
            TimeoutError: 超时无法获取连接
        """
        conn = None
        start_time = time.time()
        
        while conn is None:
            try:
                conn = self._pool.get(block=True, timeout=1)
                # 验证连接有效性
                try:
                    conn.execute("SELECT 1").fetchone()
                    return conn
                except sqlite3.Error:
                    # 连接无效，创建新连接
                    conn.close()
                    conn = None
                    with self._lock:
                        self._size -= 1
                    self._add_connection()
            except queue.Empty:
                # 队列为空，尝试创建新连接
                if self._add_connection():
                    continue
                
                # 如果超过超时时间，抛出异常
                if time.time() - start_time > self.timeout:
                    raise TimeoutError("获取数据库连接超时")
    
    def return_connection(self, conn):
        """将连接归还到连接池
        
        Args:
            conn: 数据库连接
        """
        try:
            # 检查连接是否有效
            conn.execute("SELECT 1").fetchone()
            # 回滚任何未提交的事务
            conn.rollback()
            # 归还到池中
            self._pool.put(conn)
        except sqlite3.Error:
            # 连接无效，关闭并创建新连接
            try:
                conn.close()
            except:
                pass
            with self._lock:
                self._size -= 1
            self._add_connection()
    
    def close_all(self):
        """关闭所有连接"""
        with self._lock:
            while not self._pool.empty():
                try:
                    conn = self._pool.get_nowait()
                    conn.close()
                except:
                    pass
            self._size = 0


@contextmanager
def get_db_connection():
    """数据库连接上下文管理器
    
    使用示例:
    ```
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM users")
        rows = cursor.fetchall()
    ```
    
    Yields:
        sqlite3.Connection: 数据库连接
    """
    # 获取数据库路径
    db_path = current_app.config['DB_PATH']
    
    # 获取或创建连接池
    pool = SQLiteConnectionPool(db_path)
    
    # 获取连接
    conn = pool.get_connection()
    
    try:
        # 提供连接
        yield conn
    except Exception as e:
        # 出现异常回滚
        conn.rollback()
        raise e
    else:
        # 没有异常则提交
        conn.commit()
    finally:
        # 归还连接
        pool.return_connection(conn)


class DBHelper:
    """数据库帮助类，提供常用的数据库操作方法"""
    
    @staticmethod
    def execute(query: str, params: tuple = None) -> int:
        """执行SQL语句
        
        Args:
            query: SQL查询语句
            params: 查询参数
            
        Returns:
            int: 影响的行数或最后插入的ID
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params or ())
            return cursor.lastrowid or cursor.rowcount
    
    @staticmethod
    def execute_many(query: str, params_list: List[tuple]) -> int:
        """批量执行SQL语句
        
        Args:
            query: SQL查询语句
            params_list: 参数列表
            
        Returns:
            int: 影响的行数
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.executemany(query, params_list)
            return cursor.rowcount
    
    @staticmethod
    def fetch_one(query: str, params: tuple = None) -> Optional[Dict[str, Any]]:
        """获取单条记录
        
        Args:
            query: SQL查询语句
            params: 查询参数
            
        Returns:
            Dict[str, Any]: 记录字典或None
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params or ())
            row = cursor.fetchone()
            return dict(row) if row else None
    
    @staticmethod
    def fetch_all(query: str, params: tuple = None) -> List[Dict[str, Any]]:
        """获取多条记录
        
        Args:
            query: SQL查询语句
            params: 查询参数
            
        Returns:
            List[Dict[str, Any]]: 记录字典列表
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params or ())
            rows = cursor.fetchall()
            return [dict(row) for row in rows]
    
    @staticmethod
    def fetch_value(query: str, params: tuple = None) -> Any:
        """获取单个值
        
        Args:
            query: SQL查询语句
            params: 查询参数
            
        Returns:
            Any: 查询结果的第一行第一列的值
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params or ())
            row = cursor.fetchone()
            return row[0] if row else None
    
    @staticmethod
    def count(table: str, where_clause: str = "", params: tuple = None) -> int:
        """获取记录数量
        
        Args:
            table: 表名
            where_clause: WHERE子句（不包含WHERE关键字）
            params: 查询参数
            
        Returns:
            int: 记录数量
        """
        query = f"SELECT COUNT(*) FROM {table}"
        if where_clause:
            query += f" WHERE {where_clause}"
            
        return DBHelper.fetch_value(query, params)
    
    @staticmethod
    def paginate(
        table: str,
        page: int = 1,
        size: int = 10,
        where_clause: str = "",
        order_by: str = "id DESC",
        params: tuple = None
    ) -> Tuple[List[Dict[str, Any]], int]:
        """分页查询
        
        Args:
            table: 表名
            page: 页码，从1开始
            size: 每页大小
            where_clause: WHERE子句（不包含WHERE关键字）
            order_by: ORDER BY子句
            params: 查询参数
            
        Returns:
            Tuple[List[Dict[str, Any]], int]: (记录列表, 总记录数)
        """
        # 获取总记录数
        total = DBHelper.count(table, where_clause, params)
        
        # 构建查询
        query = f"SELECT * FROM {table}"
        if where_clause:
            query += f" WHERE {where_clause}"
        
        query += f" ORDER BY {order_by} LIMIT ? OFFSET ?"
        
        # 计算偏移量
        offset = (page - 1) * size
        
        # 添加分页参数
        new_params = list(params or ())
        new_params.extend([size, offset])
        
        # 执行查询
        records = DBHelper.fetch_all(query, tuple(new_params))
        
        return records, total 